#include "DpaCommon.h"

#define PRINTF_BUFF_BSIZE (4 * 2048)

DOCA_LOG_REGISTER(WRAPPER::FLEXIO);

namespace FLEX {
    Context::Context(std::string device_name) {
        struct ibv_device **dev_list{};
        struct ibv_device *dev{};

        dev_list = ibv_get_device_list(nullptr);
        assert(dev_list != nullptr);

        for (int i = 0; dev_list[i] != nullptr; i++) {
            if (strcmp(dev_list[i]->name, device_name.c_str()) == 0) {
                dev = dev_list[i];
                break;
            }
        }
        if (dev == nullptr) {
            LOG(INFO) << "Device name not found";
            for (int i = 0; dev_list[i] != nullptr; i++) {
                printf("%s\n", dev_list[i]->name);
            }
            exit(1);
        }
        assert(dev != nullptr);

        inner_ibv_context = ibv_open_device(dev);
        assert(inner_ibv_context != nullptr);
        ibv_free_device_list(dev_list);
    }

    Context::~Context() noexcept(false) {
        int result = ibv_close_device(inner_ibv_context);
        assert(result == 0);
    }

    ::ibv_context *Context::get_inner_ptr() {
        return inner_ibv_context;
    }

    std::string Context::get_type_name() const {
        return "ibv_context";
    }


    ::flexio_uar *Context::get_flexio_uar() {
        return flexio_uar;
    }

    uint32_t Context::get_flexio_uar_id() {
        assert(flexio_uar != nullptr);
        return flexio_uar_get_id(flexio_uar);
    }

    ::ibv_pd *Context::get_pd() {
        return pd;
    }

    ::flexio_process *Context::get_process() {
        return process;
    }

    ::flexio_window *Context::get_window() {
        assert(window != nullptr);
        return window;
    }

    uint32_t Context::get_window_id() {
        assert(window != nullptr);
        uint32_t window_id = flexio_window_get_id(window);
        return window_id;
    }

    void Context::alloc_pd() {
        pd = ibv_alloc_pd(inner_ibv_context);
        assert(pd != nullptr);
    }

    void *Context::alloc_memalign(size_t size, size_t align_size) {
        void * ptr = nullptr;
        int ret = posix_memalign(&ptr, align_size, size);
        if (ret != 0) {
            LOG(INFO) << "posix_memalign failed";
            return nullptr;
        }
        return ptr;
    }

    void Context::create_process(flexio_app *app_device) {
        flexio_process_attr process_attr = { nullptr, 0, nullptr};
        // flexio_process_attr process_attr;
        flexio_status ret = flexio_process_create(inner_ibv_context, app_device, &process_attr, &process);
        assert(ret == FLEXIO_STATUS_SUCCESS);
    }

    void Context::generate_flexio_uar() {
        assert(flexio_uar == nullptr);
        flexio_uar = flexio_process_get_uar(process);
        assert(flexio_uar != nullptr);
    }

    void Context::create_window() {
        flexio_status ret = flexio_window_create(process, pd, &window);
        assert(ret == FLEXIO_STATUS_SUCCESS);
    }

    ::ibv_mr *Context::get_mr(size_t index) {
        assert(index < mr_vec.size());
        return mr_vec[index];
    }

    ibv_mr *Context::reg_mr(void *addr, size_t size, int flag) {
        /*Must be 64B aligned buffer, posix_memalign(&buf, 64, size);*/
        ibv_mr *mr = ibv_reg_mr(pd, addr, size, flag);
        assert(mr != nullptr);
        mr_vec.push_back(mr);
        return mr;
    }

    void Context::print_init(const std::string &f_name_prefix, int num_threads) {
        flexio_msg_stream_attr_t attr;
        attr.uar = flexio_uar;
        attr.data_bsize = PRINTF_BUFF_BSIZE;
        attr.sync_mode = FLEXIO_LOG_DEV_SYNC_MODE_SYNC;
        attr.level = FLEXIO_MSG_DEV_DEBUG;
        attr.stream_name = const_cast<char *>("Default Stream");
        attr.mgmt_affinity.type = FLEXIO_AFFINITY_NONE;

        // add to test
        attr.transport_mode = flexio_msg_transport::FLEXIO_MSG_TRANSPORT_QP_RC;
        //

        flexio_status ret = flexio_msg_stream_create(process, &attr, stdout, nullptr, &default_stream);
        assert(ret == FLEXIO_STATUS_SUCCESS);

        std::string prefix = "./log/" + f_name_prefix + "_";
        streams = static_cast<flexio_msg_stream **>(malloc(num_threads * sizeof(flexio_msg_stream *)));
        number_of_threads = num_threads;
        for (int i = 0; i < num_threads; i++) {
            std::string file_name = prefix + std::to_string(i) + ".txt";
            attr.stream_name = const_cast<char *>((f_name_prefix + "_" + std::to_string(i)).c_str());
            FILE *fp = fopen(file_name.c_str(), "w");
            assert(fp != nullptr);
            ret = flexio_msg_stream_create(process, &attr, fp, NULL, &streams[i]);
            assert(ret == FLEXIO_STATUS_SUCCESS);

            int id = flexio_msg_stream_get_id(streams[i]);
            assert(id == i + 1);
        }
    }

    void Context::flush() {
        flexio_msg_stream_flush(default_stream);
        for (int i = 0; i < number_of_threads; i++) {
            flexio_msg_stream_flush(streams[i]);
        }
    }

    CommandQ::CommandQ(flexio_process *process, int num_threads, int batch_size) {
        flexio_cmdq_attr cmdq_fattr;
        cmdq_fattr.workers = num_threads;
        cmdq_fattr.batch_size = batch_size;
        cmdq_fattr.state = FLEXIO_CMDQ_STATE_PENDING;

        flexio_status ret = flexio_cmdq_create(process, &cmdq_fattr, &cmd_q);
        assert(ret == FLEXIO_STATUS_SUCCESS);
    }

    CommandQ::~CommandQ() noexcept(false) {
    }

    ::flexio_cmdq *CommandQ::get_inner_ptr() {
        return cmd_q;
    }

    std::string CommandQ::get_type_name() const {
        return "flexio_cmdq";
    }

    void CommandQ::add_task(flexio_func_t func, flexio_uintptr_t ddata) {
        flexio_status ret = flexio_cmdq_task_add(cmd_q, func, ddata);
        assert(ret == FLEXIO_STATUS_SUCCESS);
    }

    void CommandQ::run() {
        flexio_status ret = flexio_cmdq_state_running(cmd_q);
        assert(ret == FLEXIO_STATUS_SUCCESS);
    }

    long CommandQ::wait_run(int time_out_seconds) { // return us
        (void)time_out_seconds;
        long duration = 0;
        while (!flexio_cmdq_is_empty(cmd_q)) {
            
        }
        return duration;
    }
}